import os
from collections.abc import Mapping

import torch

from mmdet.models import DETECTORS, SingleStageDetector


@DETECTORS.register_module
class SSDLite(SingleStageDetector):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def extract_feat(self, img):
        outs = self.backbone(img)
        if self.with_neck:
            outs += self.neck(outs[-1])
        return outs

    def init_weights(self, pretrained=None):
        if isinstance(pretrained, str) and os.path.exists(pretrained):
            state_dict = torch.load(pretrained)
            if isinstance(state_dict, Mapping) and "state_dict" in state_dict:
                state_dict = state_dict["state_dict"]
            self.load_state_dict(state_dict, strict=True)
            print(F"SSDLite load pretrained model {pretrained} succeed.")
        else:
            print(F"SSDLite load pretrained model {pretrained} failed.")
